#pragma once
#include "ApprovedExtension.h"
#include "ComBSTRResource.h"
#include "CommandTemplate.h"
#include "PEImage.h"
#include "Resource.h"

extern "C" const GUID CLSID_DllContextMenu;
extern "C" const GUID LIBID_DllContextMenuLib;

extern const CComBSTRResource ProjectName;

struct CDllContextMenuModule : CAtlDllModuleT<CDllContextMenuModule>
{
	DECLARE_LIBID(LIBID_DllContextMenuLib)

	virtual HRESULT AddCommonRGSReplacements(IRegistrarBase* pRegistrar) throw()
	{
		HRESULT hr = E_FAIL;
		LPOLESTR szClsid = NULL;

		if (FAILED(hr = StringFromCLSID(CLSID_DllContextMenu, &szClsid)))
		{
			goto hell;
		}

		if (FAILED(hr = pRegistrar->AddReplacement(OLESTR("CLSID_PEContextMenu"), szClsid)))
		{
			goto hell;
		}

		if (FAILED(hr = pRegistrar->AddReplacement(OLESTR("PROJNAME"), ProjectName)))
		{
			goto hell;
		}

hell:
		if (szClsid != NULL)
		{
			CoTaskMemFree(szClsid);
		}

		if (FAILED(hr))
		{
			return hr;
		}
		
		return CAtlDllModuleT::AddCommonRGSReplacements(pRegistrar);
	}
};

struct ATL_NO_VTABLE CDllContextMenu :
	CComObjectRootEx<CComSingleThreadModel>,
	CComCoClass<CDllContextMenu, &CLSID_DllContextMenu>,
	IShellExtInit,
	IContextMenu
{
	DECLARE_REGISTRY_RESOURCEID(IDR_PECONTEXTMENUHANDLER)
	DECLARE_NOT_AGGREGATABLE(CDllContextMenu)

	BEGIN_COM_MAP(CDllContextMenu)
		COM_INTERFACE_ENTRY(IShellExtInit)
		COM_INTERFACE_ENTRY(IContextMenu)
	END_COM_MAP()

	DECLARE_PROTECT_FINAL_CONSTRUCT()

	HRESULT FinalConstruct()
	{
		m_verbs = NULL;

		return S_OK;
	}

	void FinalRelease()
	{
		if (m_verbs != NULL)
		{
			delete m_verbs;
			m_verbs = NULL;
		}
	}

	STDMETHODIMP Initialize(LPCITEMIDLIST, LPDATAOBJECT, HKEY);
	STDMETHODIMP GetCommandString(UINT, UINT, LPUINT, LPSTR, UINT);
	STDMETHODIMP InvokeCommand(LPCMINVOKECOMMANDINFO);
	STDMETHODIMP QueryContextMenu(HMENU, UINT, UINT, UINT, UINT);

private:

	struct CVerb;
	struct CSelectedFile;
	
	typedef vector<const CVerb>         verb_vector;
	typedef vector<const CSelectedFile> item_vector;

	typedef verb_vector::const_reverse_iterator verb_iterator;
	typedef item_vector::const_iterator         item_iterator;

	item_vector  m_files;
	verb_vector* m_verbs;

	void EnsureVerbsLoaded();
	STDMETHODIMP CopyCommandString(LPVOID, UINT, BOOL, const tstring&);

	static const LPCTSTR RegPath;

	enum FileType
	{
		FileTypeInvalid = 0x0,
		FileTypeWin32   = 0x1,
		FileTypeCOM     = 0x2,
		FileTypeManaged = 0x4,

		FileTypeAny = FileTypeWin32|FileTypeCOM|FileTypeManaged,
	};

	struct CSelectedFile
	{
		explicit       CSelectedFile(const tstring& path);

		FileType       Type() const { return FileType(m_Type); }
		const tstring& Path() const { return m_Path; }
	
	private:
		tstring  m_Path;
		DWORD    m_Type;
	};

	struct CVerb
	{
		explicit CVerb(const tstring&, CRegKey&);

		bool CanInvoke(const CSelectedFile& item) const;
		bool CanInvoke(const item_vector& items) const;
		HRESULT Invoke(const item_vector& items) const;

		const tstring& Name() const
		{
			return m_Name; 
		}

		const tstring& DisplayName() const  
		{
			if (m_DisplayName.empty())
			{
				return m_Name;
			}
				
			return m_DisplayName; 
		}

		const tstring& HelpString() const
		{
			return m_HelpString;
		}

	private:

		DWORD            m_Type;
		tstring          m_Name;
		tstring          m_HelpString;
		tstring          m_DisplayName;
		CCommandTemplate m_Template;

		static tstring QueryStringValue(CRegKey&, LPCTSTR);
		static HRESULT StartProcess(const tstring&);
		static void ReportError(HRESULT);
	};
};

OBJECT_ENTRY_AUTO(CLSID_DllContextMenu, CDllContextMenu)
